package com.zhinengshidai.mcp.servers.common;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.json.JsonWriteFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zhinengshidai.mcp.servers.common.annotation.Prompt;
import com.zhinengshidai.mcp.servers.common.annotation.Resource;
import com.zhinengshidai.mcp.servers.common.annotation.Tool;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.McpSyncServerExchange;
import io.modelcontextprotocol.server.transport.HttpServletSseServerTransportProvider;
import io.modelcontextprotocol.server.transport.StdioServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.BiFunction;

/**
 * @Description
 * @Author chenss
 * @CreateTime 2025-04-01 12:29:25
 * @ModifyTime
 */
public class McpServerManager {
    private McpSyncServer syncServer;

    private String getLoggerClass(Class<?> aclass) {
        return aclass.getCanonicalName();
    }


    public void initStdioServer(String serverName, String version, Class<?> aclass,
                                McpSchema.ServerCapabilities serverCapabilities) throws Exception {
        ObjectMapper objectMapper = new ObjectMapper();
        // 客户端也得设置这两个参数，防止中文字符导致的传输错误
        objectMapper.enable(JsonWriteFeature.ESCAPE_NON_ASCII.mappedFeature());
        objectMapper.enable(JsonGenerator.Feature.ESCAPE_FORWARD_SLASHES);
        objectMapper.configure(JsonGenerator.Feature.ESCAPE_NON_ASCII, true);
        StdioServerTransportProvider transportProvider = new StdioServerTransportProvider();
        syncServer = McpServer.sync(transportProvider)
                .serverInfo(serverName, version)
                .capabilities(serverCapabilities)
                .build();

        Object newInstance = aclass.getDeclaredConstructor().newInstance();

        if (Objects.nonNull(serverCapabilities.tools())) {
            addTools(newInstance);
        }

        if (Objects.nonNull(serverCapabilities.prompts())) {
            addPrompts(newInstance);
        }

        if (Objects.nonNull(serverCapabilities.resources())) {
            addResources(newInstance);
        }

        syncServer.loggingNotification(McpSchema.
                LoggingMessageNotification.builder()
                .level(McpSchema.LoggingLevel.WARNING)
                .logger(getLoggerClass(aclass))
                .data("Server " + serverName + " initialized")
                .build());

    }


    public void initSseServer(String serverName, String version,
                              String messageEndpoint, String sseEndpoint) {
        HttpServletSseServerTransportProvider httpServletSseServerTransport =
                new HttpServletSseServerTransportProvider(new ObjectMapper(), messageEndpoint, sseEndpoint);
        syncServer = McpServer.sync(httpServletSseServerTransport)
                .serverInfo(serverName, version)
                .capabilities(McpSchema.ServerCapabilities.builder()
                        .resources(true, true)
                        .tools(true)
                        .prompts(true)
                        .logging()
                        .build())
                .build();

        syncServer.loggingNotification(McpSchema.
                LoggingMessageNotification.builder()
                .level(McpSchema.LoggingLevel.WARNING)
                .logger(getLoggerClass(this.getClass()))
                .data("Server " + serverName + " initialized")
                .build());
    }

    public void addTool(JSONObject tool,
                        BiFunction<McpSyncServerExchange, Map<String, Object>, McpSchema.CallToolResult> handler) {
        syncServer.addTool(new McpServerFeatures.SyncToolSpecification(
                new McpSchema.Tool(tool.getString("name"),
                        tool.getString("description"),
                        tool.getString("schema")), handler));
        syncServer.notifyToolsListChanged();
    }

    public void addResource(JSONObject resource, BiFunction<McpSyncServerExchange, McpSchema.ReadResourceRequest,
            McpSchema.ReadResourceResult> handler) {
        List<McpSchema.Role> audience = JSONArray.parseArray(resource.getString("audience"), McpSchema.Role.class);
        syncServer.addResource(new McpServerFeatures.SyncResourceSpecification(
                new McpSchema.Resource(resource.getString("uri"),
                        resource.getString("name"),
                        resource.getString("description"),
                        resource.getString("mimeType"),
                        new McpSchema.Annotations(audience, resource.getDouble("priority"))),
                handler
        ));
        syncServer.notifyResourcesListChanged();
    }

    public void addPrompt(JSONObject prompt, BiFunction<McpSyncServerExchange, McpSchema.GetPromptRequest,
            McpSchema.GetPromptResult> handler) {
        List<McpSchema.PromptArgument> list = JSONArray.parseArray(prompt.getString("args"))
                .stream().map(x -> {
                    JSONObject jsonObject = JSONObject.parseObject(x.toString());
                    return new McpSchema.PromptArgument(jsonObject.getString("name"),
                            jsonObject.getString("description"),
                            jsonObject.getBoolean("required"));
                }).toList();
        syncServer.addPrompt(new McpServerFeatures.SyncPromptSpecification(new McpSchema.Prompt(
                prompt.getString("name"),
                prompt.getString("description"),
                list
        ), handler));
        syncServer.notifyPromptsListChanged();
    }


    public void addTools(Object instance) {
        Method[] methods = instance.getClass().getMethods();
        for (Method method : methods) {
            Tool annotation = method.getAnnotation(Tool.class);
            if (annotation != null) {
                String schema = annotation.schema();
                String description = annotation.description();
                String name = annotation.name();
                if (name.isBlank()) {
                    name = method.getName();
                }
                JSONObject jsonObject = new JSONObject();
                jsonObject.put("name", name);
                jsonObject.put("description", description);
                jsonObject.put("schema", schema==null||schema.isBlank()?"{}":schema);
                String finalName = name;
                this.addTool(jsonObject, (exchange, stringObjectMap) -> {
                    try {
                        Object invoke = method.invoke(instance, stringObjectMap);
                        return new McpSchema.CallToolResult(List.of(new McpSchema.TextContent(invoke.toString())), false);
                    } catch (Exception e) {
                        JSONObject res = new JSONObject();
                        res.put("tool", finalName);
                        res.put("params", stringObjectMap);
                        res.put("description", "execute failed :" + JSONObject.toJSONString(e));
                        return new McpSchema.CallToolResult(List.of(new McpSchema.TextContent(res.toString())), true);
                    }
                });
            }

        }
    }

    private void addPrompts(Object instance) {
        Method[] methods = instance.getClass().getMethods();
        for (Method method : methods) {
            Prompt annotation = method.getAnnotation(Prompt.class);
            if (annotation != null) {
                JSONObject jsonObject = new JSONObject();
                String[] args = annotation.args();
                jsonObject.put("args", JSONArray.toJSONString(args));
                String description = annotation.description();
                String name = annotation.name();
                if (name.isBlank()) {
                    name = method.getName();
                }
                jsonObject.put("name", name);
                jsonObject.put("description", description);
                String finalName = name;
                this.addPrompt(jsonObject, (exchange, stringObjectMap) -> {
                    try {
                        Object invoke = method.invoke(instance, stringObjectMap);
                        JSONObject jsonObject1 = JSONObject.parseObject(invoke.toString());
                        return new McpSchema.GetPromptResult(description, List.of(new McpSchema.PromptMessage(
                                McpSchema.Role.valueOf(jsonObject1.getString("role")), new McpSchema
                                .TextContent(jsonObject1.getString("content")))));
                    } catch (Exception e) {
                        JSONObject res = new JSONObject();
                        res.put("prompt", finalName);
                        res.put("params", stringObjectMap);
                        res.put("description", "add failed");
                        res.put("exception", JSONObject.toJSONString(e));
                        return new McpSchema.GetPromptResult(description, List.of(new McpSchema.PromptMessage(
                                McpSchema.Role.USER, new McpSchema
                                .TextContent(res.toString()))));
                    }
                });
            }

        }
    }

    public void addResources(Object instance) {
        Method[] methods = instance.getClass().getMethods();
        for (Method method : methods) {
            Resource annotation = method.getAnnotation(Resource.class);
            if (annotation != null) {
                JSONObject jsonObject = new JSONObject();
                String uri = annotation.uri();
                String mimeType = annotation.mimeType();
                String description = annotation.description();
                String name = annotation.name();
                String[] audience = annotation.audience();
                double priority = annotation.priority();
                if (name.isBlank()) {
                    name = method.getName();
                }
                jsonObject.put("name", name);
                jsonObject.put("description", description);
                jsonObject.put("uri", uri);
                jsonObject.put("mimeType", mimeType);
                jsonObject.put("audience", JSONObject.toJSONString(audience));
                jsonObject.put("priority", priority);
                this.addResource(jsonObject, (exchange, readResourceRequest) ->{
                    try {
                        Object invoke = method.invoke(instance, readResourceRequest);
                        JSONObject result = JSONObject.parseObject(invoke.toString());
                        String string = result.getString("mimeType");
                        if(string.equals("text") || string.isEmpty()){
                            return new McpSchema.ReadResourceResult(List.of(new McpSchema.TextResourceContents(
                                    result.getString("uri"),
                                    result.getString("mimeType"),
                                    result.getString("text")
                            )));
                        }else {
                            return new McpSchema.ReadResourceResult(List.of(new McpSchema.BlobResourceContents(
                                    result.getString("uri"),
                                    result.getString("mimeType"),
                                    result.getString("blob")
                            )));
                        }
                    } catch (Exception e) {
                        return new McpSchema.ReadResourceResult(List.of(new McpSchema.TextResourceContents(
                                jsonObject.getString("uri"),
                                jsonObject.getString("mimeType"),
                                "read resource error :"+JSONObject.toJSONString(e)
                        )));
                    }
                });
            }

        }
    }
}
